#include <cstdint>
#include <cassert>
#include <algorithm>
#include <memory.h>
#include <cmath>
#include "Enigma.hpp"
#include "random_vector.hpp"

#define N 16    // charset_size
#define M 3     // rotors_number


/// @brief Test operator()
/// @param key 
static void test_operator(const Enigma::Key<M>& key)
{
    // Given
    clay::Enigma<N, M> enigma(key);
    std::vector<std::uint8_t> origin, origin_copy;
    origin_copy = origin = generateRandomVector<std::uint8_t, N>(0, N - 1);

    // When
    enigma(origin);
    // Then
    assert(!std::equal(origin.begin(), origin.end(), origin_copy.begin()));
    // When
    enigma(origin);
    // Then
    assert(std::equal(origin.begin(), origin.end(), origin_copy.begin()));
}


/// @brief Test the update function.
/// @param key 
static void test_update(const Enigma::Key<M>& key)
{
    // Given
    Enigma::Encryptor<N, M> encryptor(key);
    Enigma::Decryptor<N, M> decryptor(key);
    std::vector<std::uint8_t> origin, origin_1;
    std::vector<std::uint8_t> origin_copy, origin_1_copy;

    origin = origin_copy = generateRandomVector<std::uint8_t, 100>(0, N - 1);
    origin_1 = origin_1_copy = generateRandomVector<std::uint8_t, 100>(0, N - 1);

    // When
    encryptor.update(origin);
    encryptor.update(origin_1);
    // Then
    assert(!std::equal(origin.begin(), origin.end(), origin_copy.begin()));
    assert(!std::equal(origin_1.begin(), origin_1.end(), origin_1_copy.begin()));
    // When
    decryptor.update(origin);
    decryptor.update(origin_1);
    // Then
    assert(std::equal(origin.begin(), origin.end(), origin_copy.begin()));
    assert(std::equal(origin_1.begin(), origin_1.end(), origin_1_copy.begin()));
}


/// @brief Test for data whose charset_size = 256.
/// @param key 
static void test_update_256(const Enigma::Key<M>& key)
{
    // Given
    const std::uint16_t CHARSET_SIZE = 256;
    Enigma::Encryptor<CHARSET_SIZE, M> encryptor(key);
    Enigma::Decryptor<CHARSET_SIZE, M> decryptor(key);
    std::vector<std::uint8_t> origin, origin_1;
    std::vector<std::uint8_t> origin_copy, origin_1_copy;

    origin = origin_copy = generateRandomVector<std::uint8_t, 10000>(0, CHARSET_SIZE - 1);
    origin_1 = origin_1_copy = generateRandomVector<std::uint8_t, 10000>(0, CHARSET_SIZE - 1);

    // When
    encryptor.update(origin);
    encryptor.update(origin_1);
    // Then
    assert(!std::equal(origin.begin(), origin.end(), origin_copy.begin()));
    assert(!std::equal(origin_1.begin(), origin_1.end(), origin_1_copy.begin()));
    // When
    decryptor.update(origin);
    decryptor.update(origin_1);
    // Then
    assert(std::equal(origin.begin(), origin.end(), origin_copy.begin()));
    assert(std::equal(origin_1.begin(), origin_1.end(), origin_1_copy.begin()));
}


/// @brief Test for big plain old data.
/// @param key 
static void test_update_plain_old_data(const Enigma::Key<M>& key)
{
    // Given
    Enigma::Encryptor<N, M> encryptor(key);
    Enigma::Decryptor<N, M> decryptor(key);
    size_t i;
    const size_t data_size = std::pow(N, M) + 100;

    void *data = malloc(data_size);
    void *data_copy = malloc(data_size);
    std::uint8_t *data_p = static_cast<std::uint8_t *>(data);

    if (data == nullptr || data_copy == nullptr) {
        std::cerr << "No memory!" << std::endl;
        
        if (data != nullptr) free(data);
        if (data_copy != nullptr) free(data_copy);

        std::abort();
    }

    // When
    for (i = 0; i < data_size; i++) {
        data_p[i] = i % N;
    }
    memcpy(data_copy, data, data_size);

    // When
    encryptor.update(data, data_size);
    // Then
    assert(memcmp(data, data_copy, data_size) != 0);
    // When
    decryptor.update(data, data_size);
    // Then
    assert(memcmp(data, data_copy, data_size) == 0);

    free(data);
    free(data_copy);
}


/// @brief Test encrypt and decrypt data whose charset_size = 256.
/// @param key 
static void test_en_and_de_256_data(const Enigma::Key<M>& key)
{
    // Given
    const std::uint16_t CHARSET_SIZE = 256;
    Enigma::Encryptor<CHARSET_SIZE, M> encryptor(key);
    Enigma::Decryptor<CHARSET_SIZE, M> decryptor(key);
    size_t i;
    const size_t data_size = 123;

    void *data = malloc(data_size);
    void *data_copy = malloc(data_size);
    std::uint8_t *data_p = static_cast<std::uint8_t *>(data);

    if (data == nullptr || data_copy == nullptr) {
        std::cerr << "No memory!" << std::endl;
        
        if (data != nullptr) free(data);
        if (data_copy != nullptr) free(data_copy);

        std::abort();
    }

    // When
    for (i = 0; i < data_size; i++) {
        data_p[i] = i % CHARSET_SIZE;
    }
    memcpy(data_copy, data, data_size);

    // When
    encryptor.update(data, data_size);
    // Then
    assert(memcmp(data, data_copy, data_size) != 0);
    // When
    decryptor.update(data, data_size);
    // Then
    assert(memcmp(data, data_copy, data_size) == 0);

    free(data);
    free(data_copy);
}


int main()
{
    Enigma::Key<M> key = {
        .rotor_shifts = {1, 2, 3},
        .connect_pair_list = {{1, 2}, {3, 5}, {13, 7}}
    };

    test_operator(key);
    test_update(key);
    test_update_256(key);
    test_update_plain_old_data(key);
    test_en_and_de_256_data(key);

    return 0;
}